from tflearn.layers.conv import conv_2d, max_pool_2d, conv_2d_transpose
import tensorflow as tf


class FCN(object):
    def __init__(self, batch_size, img_size):
        self.batch_size = batch_size

        self.input_imgs = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, img_size, img_size])
        self.label = tf.placeholder(dtype=tf.float32, shape=[self.batch_size, img_size, img_size])

        with tf.variable_scope('Conv'):
            net = conv_2d(tf.expand_dims(self.input_imgs, -1), 16, 3, activation='relu')
            net = conv_2d(net, 32, 3, activation='relu')

        with tf.variable_scope('Deconv'):
            net = conv_2d_transpose(net, 16, 32, output_shape=[self.batch_size, img_size, img_size, 16], strides=1,
                                    bias=True)
            net = conv_2d_transpose(net, 1, 16, output_shape=[self.batch_size, img_size, img_size, 1], strides=1,
                                    bias=False,
                                    activation='sigmoid')
            output = tf.squeeze(net)
            self.den = output * 255.

        self.loss = tf.reduce_sum(tf.square(output - self.label))
        with tf.device('/cpu:0'):
            tf.scalar_summary("loss", self.loss)
            self.merged = tf.merge_all_summaries()
        tvars = tf.trainable_variables()
        grads, _ = tf.clip_by_global_norm(tf.gradients(self.loss, tvars), 5.0)
        optimizer = tf.train.MomentumOptimizer(0.01, 0.9)
        self.train_op = optimizer.apply_gradients(zip(grads, tvars))
